#!/usr/bin/env python
"""
CTPRPCClient

Client implementation of the CTP-RPC protocol.

Steps for using client object:
	1. Create instance of the client:
		cli = CTPRPCClient ( ( "127.0.0.1", 8000 ) [, "john"] [, "pippo"] )
		
		the first parameters specifies the IP address and port to which the client have to connect
		the (optional) second and third parameters are username and password for authetication

	2. Simply call published server methods on the client instance:
		ret = cli.server_method ( ... )

	3. Destroy the client to close the connection with the server:
		destroy_client ( cli )

	NOTES:
		a. The client instance has an attribute named 'method_lookup' initially set to 1.
		   If set to 0 then the client will always send the request for a method call to the server,
		   even if the method is not in the list of published methods. This feature can be used, for example,
		   if the server registers a class that supports the '__getattr__' functionality.
		   
		b. The client instance has also some other useful attributes:
		   - encrypt_alg_id		id of the encryption algorithm nogotiated with the server
		   - compress_alg_id		id of the compression algorithm nogotiated with the server
		   - client_address		( ip_address, port ) tuple of the client
		   - server_address		( ip_address, port ) tuple of the server
		   
		c. This module exports other functions:
		   - enable_encryption()		enables/disables encryption
		   - enable_compression()		enables/disables compression
		   - enable_data_compression()		enables/disables compression of extra-data (usually files)

		   - get_last_exception()		returns a list of tuples representing the server traceback of the last
		   					exception of a server method call
		   - dump_last_exception()		returns a list of strings representing the server traceback of the last
		   					exception of a server method call
		   - print_last_exception()		prints the strings returned by dump_last_exception()
"""

import sys, os, os.path
import socket
import cPickle
import traceback
import md5
import sha
import hmac
import time
import os3.utils.thread_control

from CTPRPCGlobal import *


# *******************************************************
# CTPRPCMethod class
class CTPRPCMethod ( object ):
	"""
	Questa classe e' un guscio attorno alla chiamata di un metodo.
	Viene restituita dalla __getattr__ sull'oggetto CTPRPCClient.
	"""

	def __init__ ( self, conn, method_name ):
		"""
		Questa classe viene inizializzata con un
		oggetto CTPRPCClientConnection e con il nome del metodo
		da invocare.
		"""
		
		super ( CTPRPCMethod, self ).__init__()
		self._conn = conn
		self._method_name = method_name


	# *******************************************************
	# __getattr__() - permette il recupero di un metodo 'puntato'
	def __getattr__ ( self, name ):
		return CTPRPCMethod ( self._conn, '%s.%s' % ( self._method_name, name ) )


	# *******************************************************
	# __call__() - effettua la chiamata del metodo remoto tramite CTP-RPC
	def __call__ ( self, *pos_args, **named_args ):
		"""
		Quando un istanza di questa classe viene chiamata,
		allora viene usata la connessione per inviare al
		server un comando CMD_INVOKE che ha i seguenti argomenti:
		- nome del netodo da invocare
		- lista di parametri posizionali
		- dizionario di parametri con nome

		Dopo aver inviata il comando viene letta la risposta che puo' essere:
		- RESP_OK se e' andato tutto bene; l'unico argomento contiene il valore di ritorno della funzione remota
		- RESP_BADMETHOD se il metodo invocato non esiste
		- RESP_EXCEPTION se il metodo remoto ha provocato un'eccezione; gli argomenti in questo caso saranno:
		  - classe o istanza dell'eccezione
		  - istanza o parametro dell'eccezione
		  - lista di linee di traceback
		"""
		
		self._conn._tc.lock_acquire ( "hb_lock" )
		try:
			self._conn.send_command ( CMD_INVOKE, ( self._method_name, pos_args, named_args ) )
			ret = self._conn.get_response()
		finally:
			self._conn._tc.lock_release ( "hb_lock" )

		if not ret: raise InvalidResponseError

		resp, args = ret

		if resp == RESP_OK:
			#call successfull
			return args[0]
		elif resp == RESP_EXCEPTION:
			#exception occurred, re-raise it
			self._conn._last_exception = args
			raise args[0], args[1]
		elif resp == RESP_BADMETHOD:
			#bad method, raise exception
			raise BadMethodError, "method '%s' not found" % self._method_name
		else:
			raise InvalidResponseError

	
# *******************************************************
# CTPRPCClientConnection class
class CTPRPCClientConnection ( object ):
	"""
	Classe che rappresenta una connessione con il server. Viene istanziata per
	ogni oggetto CTPRPCClient contiene tutti i metodi di comunicazione con il server.
	E' stata creata una classe a parte per lasciare l'oggetto client il piu' "pulito"
	possibile (senza funzioni che potrebbero avere lo stesso nome di funzioni remote).
	"""

	def __init__ ( self, sock ):
		"""
		La classe viene inizializzata con il socket della connessione al server.
		"""
		
		super ( CTPRPCClientConnection, self ).__init__()
		self._sock = sock
		self._handshake_complete = 0
		self._last_exception = ( None, None, [] )

		self._tc = os3.utils.thread_control.ThreadControl()

		self._encrypt_alg_id = 0
		self._compress_alg_id = 0
		self._session_id = 0
		self._seq_num = 0

		self._encrypt = 1
		self._compress = 1
		self._compress_data = 1

		self._server_methods = []


	# *******************************************************
	# close() - chiude la connessione
	def close ( self ):
		"""
		close ( self )

		Chiude la connessione al server.
		"""
		
		self._tc.quit_all()

		if self._sock:
			self._sock.shutdown(2)
			self._sock.close()
			self._sock = None
		
		
	# *******************************************************
	# handshake() - effettua lo scambio di dati iniziale per concordare
	#               gli algoritmi di compressione e cifratura
	def handshake ( self ):
		"""
		handshake ( self )

		Effettua lo scambio di dati iniziale tra client e server. Vedi la funzione analoga sul server
		per una documentazione piu' dettagliata di questa procedura.
		"""

		enc_types = supported_encryption_types.keys()
		enc_types.sort()
		enc_types.reverse()
		comp_types = supported_compression_types.keys()
		comp_types.sort()
		comp_types.reverse()

		params = [ ( 1, 0 ) ]			# version
		params.append ( tuple ( enc_types ) )	# supported encryption
		params.append ( tuple ( comp_types ) )	# supported compression
		
		if is_encryption_supported:
			#encryption is supported
			RANDPOOL.stir()
			self._client_random = RANDPOOL.get_bytes ( 32 )
			self._client_rsa_key = RSA.generate ( RSA_KEY_LENGTH, RANDPOOL.get_bytes )

			params.append ( self._client_random )					# client random
			params.append ( ( self._client_rsa_key.n, self._client_rsa_key.e ) )	# client public key

		self.send_command ( CMD_HANDSHAKE_1, tuple ( params ) )

		ret = self.get_response()
		if not ret: raise InvalidResponseError

		resp, args = ret

		if resp != RESP_HANDSHAKE_1: raise InvalidResponseError

		if len ( args ) < 4: raise InvalidResponseError
		if not isinstance ( args[0], tuple ) or not isinstance ( args[1], int ) or not isinstance ( args[2], int ): raise InvalidResponseError

		version, self._encrypt_alg_id, self._compress_alg_id, self._session_id = args[:4]

		if len ( version ) != 2 or self._encrypt_alg_id not in supported_encryption_types.keys() or self._compress_alg_id not in supported_compression_types.keys():
			raise InvalidResponseError

		self._encrypt_alg = supported_encryption_types[self._encrypt_alg_id][1]
		self._compress_alg = supported_compression_types[self._compress_alg_id][1]

		if is_encryption_supported and self._encrypt_alg and len ( args ) == 7:

			if not isinstance ( args[4], str ) or not isinstance ( args[5], tuple ) or not isinstance ( args[6], str ): raise InvalidResponseError

			self._server_random, serv_pub_key, serv_pre_key = args[4:]

			if len ( self._server_random ) != 32 or len ( serv_pub_key ) != 2: raise InvalidResponseError

			self._server_rsa_key = RSA.construct ( serv_pub_key )
			
			server_pre_key = self._client_rsa_key.decrypt ( serv_pre_key )

			#and now derive other data
			client_pre_key = RANDPOOL.get_bytes ( self._encrypt_alg.key_size )

			hmac_master_key = md5.new ( self._client_random + self._server_random + client_pre_key ).digest()
			hmac_master_key += sha.new ( self._server_random + self._client_random + server_pre_key ).digest()

			random_data = CTP_PRF ( hmac_master_key, "MASTER" + self._server_random + self._client_random, 32 + 2 * self._encrypt_alg.block_size )

			client_init_vector = random_data[:self._encrypt_alg.block_size]
			server_init_vector = random_data[self._encrypt_alg.block_size:2 * self._encrypt_alg.block_size]
			self._hmac_key = random_data[2 * self._encrypt_alg.block_size:]

			client_key = CTP_PRF ( hmac_master_key + client_pre_key, "CLIENT" + self._client_random + self._server_random, self._encrypt_alg.key_size )
			server_key = CTP_PRF ( hmac_master_key + server_pre_key, "SERVER" + self._client_random + self._server_random, self._encrypt_alg.key_size )
			
			self._client_sess_key = self._encrypt_alg.new ( client_key, self._encrypt_alg.MODE_CBC, client_init_vector )
			self._server_sess_key = self._encrypt_alg.new ( server_key, self._encrypt_alg.MODE_CBC, server_init_vector )

			#OK! adesso preparo la richiesta handshake fase 2
			enc_key = self._server_rsa_key.encrypt ( client_pre_key, 0 )[0]
			params = ( enc_key, )

			self.send_command ( CMD_HANDSHAKE_2, tuple ( params ) )

			self._handshake_complete = 1
		
			ret = self.get_response()
			if not ret: raise InvalidResponseError

			resp, args = ret

			if resp != RESP_HANDSHAKE_2: raise InvalidResponseError

			if len ( args ) < 1: raise InvalidResponseError
			if not isinstance ( args[0], str ): raise InvalidResponseError

			resp = hmac.new ( self._hmac_key, client_pre_key + server_pre_key, md5 ).digest()
			resp += hmac.new ( self._hmac_key, server_pre_key + client_pre_key, sha ).digest()

			if resp != args[0]: raise InvalidResponseError

		return 1


	# *******************************************************
	# authenticate() - effettua l'autenticazione dell'utente
	def authenticate ( self, login, passwd ):
		"""
		authenticate ( self, login, passwd )

		Effettua l'autenticazione dell'utente inviando un comando CMD_LOGIN con login e password come parametri.
		"""

		self.send_command ( CMD_LOGIN, ( login, passwd ) )
		
		ret = self.get_response()
		if not ret: raise InvalidResponseError

		resp, args = ret

		if resp != RESP_OK: raise BadLoginOrPassword

		if len(args) < 1: raise InvalidResponseError
		if not isinstance ( args[0], tuple ): raise InvalidResponseError

		self._server_methods = args[0]
	
		return 1


	# *******************************************************
	# _start_heartbeat() - avvia il thread di heartbeat
	def _start_heartbeat ( self ):
		"""
		_start_heartbeat ( self )

		Avvia il thread che si occupera di inviare i pacchetti di heartbeat.
		"""
		
		self._tc.spawn ( "hb_thread", self._heartbeat, () )


	# *******************************************************
	# _heartbeat() - invia un heartbeat al server
	def _heartbeat ( self ):
		"""
		_heartbeat ( self )

		Funzione eseguita in un thread a parte che invia ogni 60 secondi un comando
		CMD_HEARTBEAT al server per mantenere la connessione attiva.
		"""
		
		count = 0
		while self._tc.keep_run ( "hb_thread" ):
			time.sleep ( 1 )
			count += 1
			if count % 60 == 0:
				count = 0
				self._tc.lock_acquire ( "hb_lock" )
				try:
					self.send_command ( CMD_HEARTBEAT )
					ret = self.get_response()
				except:
					pass
				self._tc.lock_release ( "hb_lock" )


	# *******************************************************
	# send_command() - invia un comando
	def send_command ( self, cmd, args = None ):
		"""
		send_command ( self, cmd, args )
		
		Invia un comando al server. L'ID del comando e' specificato in ``cmd`` mentre
		``args`` e' una tupla di eventuali argomenti.
		"""
		
		compress_mod = None
		encrypt_key = None
		hmac_key = 'CTP-RPC'
		block_size = 1

		if self._handshake_complete:
			if self._compress: compress_mod = self._compress_alg
			if self._encrypt and self._encrypt_alg:
				encrypt_key = self._server_sess_key
				block_size = self._encrypt_alg.block_size
			if is_encryption_supported: hmac_key = self._hmac_key

		if not args: args = ()
		send_object ( self._sock, ( cmd, args, self._seq_num ), self._session_id,
			hmac_key, compress_mod, encrypt_key, block_size, self._compress_data )

		self._seq_num += 1

		
	# *******************************************************
	# get_response() - riceve una risposta
	def get_response ( self ):
		"""
		get_response ( self ) -> ( int, tuple )

		Riceve una risposta dal server e restituisce una tuple con l'ID della risposta
		e una tupla di eventuali parametri.
		"""

		compress_mod = None
		encrypt_key = None
		hmac_key = 'CTP-RPC'

		if self._handshake_complete:
			#print "enc_alg", self._encrypt_alg, "  enc_key: ", self._server_sess_key
			compress_mod = self._compress_alg
			if self._encrypt_alg: encrypt_key = self._server_sess_key
			if is_encryption_supported: hmac_key = self._hmac_key
	
		ret = read_object ( self._sock, self._session_id, hmac_key, compress_mod, encrypt_key )
		if not ret: return None

		obj = ret[0]
		if not isinstance ( obj, tuple ) or len ( obj ) != 3: return None

		resp, args, seq_num = obj
		if not isinstance ( resp, int ) or not isinstance ( args, tuple ): return None

		if seq_num != self._seq_num: return None
		self._seq_num += 1

		return resp, args


# *******************************************************
# CTPRPCClient class
class CTPRPCClient ( object ):
	"""
	Classe principale sulla quale chiamare i metodi che verranno eseguiti in remoto.
	"""

	def __init__ ( self, connect_addr, login = '', passwd = '', delay_connect = 0, socket_timeout = 300 ):
		"""
		Questa classe viene inizializzata con l'indirizzo del server, la login e la password di connessione.
		L'indirizzo del server e' una tupla che contiene indirizzo IP e porta.
		"""

		super ( CTPRPCClient, self ).__init__()

		self.method_lookup = 1
		self._ctp_connection = None

		self._connect_params = ( connect_addr, login, passwd, socket_timeout )
		self._sock = None

		if not delay_connect: _client_connect ( self )


	def __del__ ( self ):
		"""
		E' necessario cancellare esplicitamente un oggetto client per chiudere la connessione al server.
		"""
		
		if self._ctp_connection: self._ctp_connection.close()


	# *******************************************************
	# __getattr__() - chiamata per ogni richiesta di funzione
	def __getattr__ ( self, attr ):
		"""
		Guarda se il metodo cercato e' presente nella lista dei metodi inviata dal server (questo solo se
		self.method_lookup e' 1). Dopodiche' crea un'istanza che rappresenta il metodo e la restituisce.
		"""

		if not self._sock: _client_connect ( self )
		
		if self.method_lookup and not attr in self._ctp_connection._server_methods: raise AttributeError, attr
		return CTPRPCMethod ( self._ctp_connection, attr )


# *******************************************************
# _client_connect () - apre la connessione
def _client_connect ( client ):
	client._sock = socket.socket ( socket.AF_INET, socket.SOCK_STREAM, 0 )

	#connect
	client._sock.connect ( client._connect_params[0] )
	client._sock.settimeout ( client._connect_params[3] )		# 5 minutes by default

	client.client_address = client._sock.getsockname()
	client.server_address = client._sock.getpeername()

	#create a new connection
	client._ctp_connection = CTPRPCClientConnection ( client._sock )

	client._ctp_connection.handshake()
	client._ctp_connection.authenticate ( client._connect_params[1], client._connect_params[2] )

	client._ctp_connection._start_heartbeat()

	client.encrypt_alg_id = client._ctp_connection._encrypt_alg_id
	client.compress_alg_id = client._ctp_connection._compress_alg_id

	del client._connect_params

	
# *******************************************************
# destroy_client() - chiude la connessione
def destroy_client ( client ):
	"""
	destroy_client ( client )
	
	chiude la connessione con il server.
	"""

	client.__del__()


# *******************************************************
# get_last_exception() - restituisce l'ultima eccezione che si e' verificata dopo una chiamata di procedura remota.
def get_last_exception ( client ):
	"""
	get_last_exception ( client )
	
	restituisce l'ultima eccezione che si e' verificata dopo una chiamata di procedura remota.
	"""
	
	return client._ctp_connection._last_exception

	
# *******************************************************
# dump_last_exception() - resituisce una lista di stringhe che contengono
#                         il traceback dell'ultima eccezione
def dump_last_exception ( client ):
	"""
	dump_last_exception ( client )

	resituisce una lista di stringhe che contengono il traceback dell'ultima eccezione
	"""
	
	retval = [ "Traceback (most recent call last):\n" ]
	retval += traceback.format_list ( client._ctp_connection._last_exception[2] )
	retval += traceback.format_exception_only ( client._ctp_connection._last_exception[0], client._ctp_connection._last_exception[1] )
	return retval


# *******************************************************
# print_last_exception() - stampa a video il traceback dell'ultima eccezione
def print_last_exception ( client ):
	"""
	print_last_exception ( client )

	stampa a video il traceback dell'ultima eccezione
	"""
	
	tb = dump_last_exception ( client )
	for l in tb:
		print l.rstrip()


# *******************************************************
# enable_encryption() - abilita/disabilita la crittografia
def enable_encryption ( client, enabled = 1 ):
	"""
	enable_encryption ( client, enabled )

	abilita/disabilita la crittatura dei dati
	"""
	client._ctp_connection._encrypt = enabled


# *******************************************************
# enable_compression() - abilita/disabilita la compressione
def enable_compression ( client, enabled = 1 ):
	"""
	enable_compression ( client, enabled )

	abilita/disabilita la compressione dei dati
	"""
	client._ctp_connection._compress = enabled


# *******************************************************
# enable_data_compression() - abilita/disabilita la compressione dei dati extra (files, ...)
def enable_data_compression ( client, enabled = 1 ):
	"""
	enable_compression ( client, enabled )

	abilita/disabilita la compressione dei dati extra (files, ...)
	"""
	client._ctp_connection._compress_data = enabled


######################################################
###                    TEST AREA                   ###
######################################################
if __name__ == '__main__':
	def test():
		cli = CTPRPCClient ( ( "127.0.0.1", 8000 ) )
	
		try:
			try:
				import tempfile
				try:
					ret = cli.test ( "ciao", file ( "CTPRPCServer.py", "rb" ) )
				except InvalidResponseError:
					print "connessione chiusa"
					sys.exit(0)

				#i dati della connessione sono in cli._ctp_connection
				print "encrypt_alg: ", cli.encrypt_alg_id
				print "compress_alg:", cli.compress_alg_id
				print "client_addr: ", cli.client_address
				print "server_addr: ", cli.server_address

				print ret
				
			except:
				raise
				print_last_exception ( cli )
		finally:
			destroy_client ( cli )

	test()


